#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
功能超级完善的点云去噪脚本 (A Super Comprehensive Point Cloud Denoising Script)
作者: [你的名字或AI助手名称]
版本: 1.0
------------------------------------------------------------------------------------
本脚本提供了一个统一的命令行界面，用于对点云文件进行去噪处理。
集成了多种主流的点云去噪算法，并允许用户通过参数进行灵活选择和调整。

核心功能:
1. 支持多种去噪算法:
   - 统计异常值移除 (Statistical Outlier Removal)
   - 半径异常值移除 (Radius Outlier Removal)
   - 体素下采样去噪 (Voxel Downsampling)
   - 基于聚类的去噪 (DBSCAN Clustering)
2. 灵活的参数调整，适用于不同场景。
3. 可选的去噪前后效果可视化对比。
4. 可选的噪声点云保存功能，便于分析。
5. 清晰的性能报告（处理时间、移除点数等）。

依赖库:
- open3d
- numpy

使用示例:
# 使用统计异常值移除 (SOR)
python denoise_point_cloud.py input.pcd output_sor.pcd --method sor --nb_neighbors 20 --std_ratio 1.0

# 使用半径异常值移除 (ROR) 并可视化
python denoise_point_cloud.py input.ply output_ror.ply --method ror --nb_points 10 --radius 0.05 --visualize

# 使用DBSCAN去噪，并保存噪声点
python denoise_point_cloud.py input.xyz output_dbscan.xyz --method dbscan --eps 0.02 --min_points 10 --save_noise
"""

import open3d as o3d
import numpy as np
import argparse
import time
import os

def visualize_comparison(pcd_original, pcd_denoised):
    """
    可视化去噪前后的点云对比。
    原始点云为灰色，去噪后点云为红色。
    """
    print("INFO: 正在准备可视化窗口...")
    # 为点云上色以便区分
    pcd_original.paint_uniform_color([0.5, 0.5, 0.5])  # 灰色
    pcd_denoised.paint_uniform_color([0.9, 0.1, 0.1])  # 红色

    # 创建一个窗口并添加点云
    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name="去噪前后对比 (灰:原始, 红:去噪后)", width=1280, height=720)
    
    # 添加几何体
    vis.add_geometry(pcd_original)
    vis.add_geometry(pcd_denoised)
    
    # 设置渲染选项
    opt = vis.get_render_option()
    opt.background_color = np.asarray([0.1, 0.1, 0.1]) # 深色背景
    opt.point_size = 1.5
    
    print("INFO: 按 'q' 关闭可视化窗口。")
    vis.run()
    vis.destroy_window()

def main():
    parser = argparse.ArgumentParser(
        description="一个功能全面的点云去噪脚本，集成了多种算法。",
        formatter_class=argparse.RawTextHelpFormatter  # 保持帮助信息中的格式
    )
    
    # --- 基本参数 ---
    parser.add_argument("input_file", help="输入的点云文件路径 (支持 .pcd, .ply, .xyz 等)。")
    parser.add_argument("output_file", help="输出的去噪后点云文件路径。")
    parser.add_argument(
        "--method",
        type=str,
        required=True,
        choices=['sor', 'ror', 'voxel', 'dbscan'],
        help="""选择去噪算法:
  sor     - 统计异常值移除 (Statistical Outlier Removal)。
            基于点与其邻域的距离分布来移除异常点。
            适用于移除稀疏的、孤立的噪声点。
  ror     - 半径异常值移除 (Radius Outlier Removal)。
            移除在指定半径内邻居数量过少的点。
            适用于控制点云的密度，移除密度低的区域。
  voxel   - 体素下采样 (Voxel Grid Downsampling)。
            通过将点云划分到体素网格并平均每个体素内的点来实现下采样。
            这不仅能简化点云，还能有效平滑表面，去除高频噪声。
  dbscan  - 基于密度的空间聚类 (DBSCAN)。
            将点云聚类，并将不属于任何簇的小规模点群视为噪声移除。
            功能强大，对噪声团块和孤立点都有效。
"""
    )

    # --- 算法特定参数 ---
    sor_group = parser.add_argument_group('SOR 参数 (当 --method=sor 时使用)')
    sor_group.add_argument('--nb_neighbors', type=int, default=20, help='SOR: 用于计算平均距离的邻居数量。')
    sor_group.add_argument('--std_ratio', type=float, default=1.0, help='SOR: 标准差乘数。值越小，去噪越保守；值越大，去噪越激进。')

    ror_group = parser.add_argument_group('ROR 参数 (当 --method=ror 时使用)')
    ror_group.add_argument('--nb_points', type=int, default=16, help='ROR: 在半径内要求的最小邻居数量。')
    ror_group.add_argument('--radius', type=float, default=0.05, help='ROR: 定义邻域的半径。')
    
    voxel_group = parser.add_argument_group('Voxel 参数 (当 --method=voxel 时使用)')
    voxel_group.add_argument('--voxel_size', type=float, default=0.01, help='Voxel: 体素的大小。值越大，点云越稀疏，平滑效果越强。')

    dbscan_group = parser.add_argument_group('DBSCAN 参数 (当 --method=dbscan 时使用)')
    dbscan_group.add_argument('--eps', type=float, default=0.02, help='DBSCAN: 定义邻域的距离(epsilon)。')
    dbscan_group.add_argument('--min_points', type=int, default=10, help='DBSCAN: 形成核心点所需的最小邻居数。')

    # --- 其他功能参数 ---
    other_group = parser.add_argument_group('其他选项')
    other_group.add_argument('--visualize', action='store_true', help='处理后，显示原始点云和去噪后点云的对比。')
    other_group.add_argument('--save_noise', action='store_true', help='将识别出的噪声点保存到单独的文件中 (文件名将在输出文件后加 "_noise")。')

    args = parser.parse_args()

    # --- 1. 加载点云 ---
    print(f"INFO: 正在从 '{args.input_file}' 加载点云...")
    try:
        pcd = o3d.io.read_point_cloud(args.input_file)
        if not pcd.has_points():
            raise ValueError("错误: 文件加载成功，但点云为空。")
    except Exception as e:
        print(f"ERROR: 加载文件失败: {e}")
        return
    
    points_before = len(pcd.points)
    print(f"INFO: 加载成功。原始点云包含 {points_before} 个点。")
    
    # 复制一份原始点云用于可视化
    if args.visualize:
        pcd_original_for_vis = o3d.geometry.PointCloud(pcd)

    # --- 2. 执行去噪 ---
    start_time = time.time()
    denoised_pcd = None
    noise_pcd = None

    print(f"\nINFO: 使用 '{args.method.upper()}' 算法进行去噪...")

    if args.method == 'sor':
        print(f"  > 参数: nb_neighbors={args.nb_neighbors}, std_ratio={args.std_ratio}")
        denoised_pcd, ind = pcd.remove_statistical_outlier(
            nb_neighbors=args.nb_neighbors,
            std_ratio=args.std_ratio
        )
        if args.save_noise:
            noise_pcd = pcd.select_by_index(ind, invert=True)

    elif args.method == 'ror':
        print(f"  > 参数: nb_points={args.nb_points}, radius={args.radius}")
        denoised_pcd, ind = pcd.remove_radius_outlier(
            nb_points=args.nb_points,
            radius=args.radius
        )
        if args.save_noise:
            noise_pcd = pcd.select_by_index(ind, invert=True)

    elif args.method == 'voxel':
        print(f"  > 参数: voxel_size={args.voxel_size}")
        denoised_pcd = pcd.voxel_down_sample(voxel_size=args.voxel_size)
        # 体素下采样是重构点云，无法直接分离噪声点，因此--save_noise在此模式下无效
        if args.save_noise:
            print("WARNING: 'voxel' 模式不支持保存噪声点，该选项将被忽略。")
    
    elif args.method == 'dbscan':
        print(f"  > 参数: eps={args.eps}, min_points={args.min_points}")
        labels = np.array(pcd.cluster_dbscan(
            eps=args.eps,
            min_points=args.min_points,
            print_progress=True
        ))
        
        # 标签为-1的点是噪声
        noise_indices = np.where(labels == -1)[0]
        inlier_indices = np.where(labels != -1)[0]
        
        denoised_pcd = pcd.select_by_index(inlier_indices)
        if args.save_noise:
            noise_pcd = pcd.select_by_index(noise_indices)

    end_time = time.time()
    
    # --- 3. 报告结果 ---
    points_after = len(denoised_pcd.points)
    points_removed = points_before - points_after
    
    print("\n----------- 去噪结果报告 -----------")
    print(f"算法          : {args.method.upper()}")
    print(f"处理时间      : {end_time - start_time:.4f} 秒")
    print(f"原始点数      : {points_before}")
    print(f"去噪后点数    : {points_after}")
    print(f"移除点数      : {points_removed} ({points_removed/points_before:.2%})")
    print("------------------------------------")

    # --- 4. 保存结果 ---
    print(f"INFO: 正在将去噪后的点云保存到 '{args.output_file}'...")
    o3d.io.write_point_cloud(args.output_file, denoised_pcd)
    print("INFO: 保存成功。")

    if args.save_noise and noise_pcd is not None:
        noise_output_path = f"{os.path.splitext(args.output_file)[0]}_noise{os.path.splitext(args.output_file)[1]}"
        print(f"INFO: 正在将噪声点云保存到 '{noise_output_path}'...")
        o3d.io.write_point_cloud(noise_output_path, noise_pcd)
        print("INFO: 噪声点云保存成功。")

    # --- 5. 可视化 ---
    if args.visualize:
        visualize_comparison(pcd_original_for_vis, denoised_pcd)
        
    print("\nINFO: 脚本执行完毕。")

if __name__ == '__main__':
    main()